
"""
情节设计评分
"""
from models import Novel
from pydantic import Field
from core.utils import BaseAnalyzer, prompt_dir


class PlotRating(BaseAnalyzer):

    def __init__(self):
        super().__init__(prompt_file=prompt_dir / 'rating/r1. 情节设计评分.md')

    def invoke(
            self,
            novel: Novel = Field(description="小说")
    ):
        all_plots = {
            chapter.chapter_title: chapter.plot for chapter in novel.chapters
        }
        plot_problems = dict()
        for chapter in novel.chapters:
            if not isinstance(chapter.problems, list):
                plot_problems[chapter.chapter_title] = chapter.problems
                continue
            plot_problems[chapter.chapter_title] = []
            for problem in chapter.problems:
                if not isinstance(problem, dict):
                    plot_problems[chapter.chapter_title].append(problem)
                elif 'dimension'  not in problem:
                    plot_problems[chapter.chapter_title].append(problem)
                elif problem['dimension'] == '情节设计':
                    plot_problems[chapter.chapter_title].append(problem)

        prompt = self.template_prompt.format(
            inputs={
                "topn": novel.topn,
                "plot": all_plots,
                "plot_problems": plot_problems
            },
            remove_template_variables=True
        )
        novel.plot_rating = self.analyze(prompt)
        return novel.plot_rating
